#include "nmqj_alg.hpp"

nmqj_alg::nmqj_alg():effective_ensemble(),channels(),MyRandom(),hamilton(),
		     initial_state(),transition_map() {}
nmqj_alg::nmqj_alg(double T,double dt,int M,effective_ensemble eff,
		   channels chans,
		   MyRandom rng,hamilton H,initial_state psi0,
		   transition_map tm):
  effective_ensemble(eff),channels(chans),MyRandom(rng),hamilton(H),
  initial_state(psi0),transition_map(tm),T_(T),dt_(dt),M_(M) {}

void nmqj_alg::det_evol(double t, double dt){
  vector<std::complex<double> > vec;
  for(int k=0;k<effective_ensemble::get_number_of_states();k++){
    vec=effective_ensemble::get_state(k);
    
    vec=vec+std::complex<double>(0.0,-dt)*(prod(hamilton::get_op(),vec)-
     					   std::complex<double>(0.0,0.5)*
     					   prod(channels::get_lamb_shift(t),vec));
    vec=vec/norm_2(vec);
    effective_ensemble::set_state(k,vec);
  }
}





void nmqj_alg::make_jumps(){
  
  //Calculate jumps
  double ps=sum(transition_probs);
  std::map<int,int> n_jumps;
  for(int ii=0;ii<effective_ensemble::occupancy(source_state);ii++){
    
    if(MyRandom::sample()<=ps){
      //hyppy
      if(target_states.size()==1){
	n_jumps[target_states(0)] = n_jumps[target_states(0)]+1;
      }
      else{
	int targ=target_states(MyRandom::sample_int_from(transition_probs));
	n_jumps[targ] = n_jumps[targ]+1;
      }
    }
  }
  //Update ensemble
  std::map<int,int>::iterator it;
  int ntot=0;
  for(it=n_jumps.begin();it!=n_jumps.end();it++){
    ntot += it->second;
    effective_ensemble::set_occupancy(it->first,
				      effective_ensemble::occupancy(it->first)+it->second);
  }
  effective_ensemble::set_occupancy(source_state,
				    effective_ensemble::occupancy(source_state)-ntot);
}

void nmqj_alg::jumps(double t,double dt){
  int n_s=effective_ensemble::get_number_of_states();
  int n_c=channels::get_number_of_channels();
  
  target_states.resize(n_c*n_s);
  transition_probs.resize(n_c*n_s);
  
  left_const_iterator il,ile;
  right_const_iterator ir,ire;
  
  std::map<int,int>::iterator ens_it;
  
  //Iterate over the states
  for(ens_it=effective_ensemble::occ.begin();ens_it!=effective_ensemble::occ.end();ens_it++){
    
    //Set source state
    source_state=ens_it->first;
    
    int num_of_transition=0;
    
    for(int chan_idx=0;chan_idx<n_c;chan_idx++){//iterate over decay channels
      
      //Get decay rate for channel chan_idx
      rate delta=get_decay_rate(chan_idx);
      
      if(delta(t)>=0){//Jumps for positive channel chan_idx
	
	lpair targs=transition_map::get_map(chan_idx).left.equal_range(ens_it->first);
	
	if(targs.first!=transition_map::get_map(chan_idx).left.end()){
	  
	  for(left_const_iterator il=targs.first;il!=targs.second;il++){
	    
	    transition_probs(num_of_transition)=
	      pos_pjump(t,dt,chan_idx,ens_it);
	    
	    target_states(num_of_transition)=il->second;
	    
	    num_of_transition += 1;
	  }
	}
      }
     
      
      //negative channel
      else if(delta(t)<0){
	//find if jump away from state hh via channel ii is possible
	rpair targs=transition_map::get_map(chan_idx).right.equal_range(ens_it->first);
	if(targs.first!=transition_map::get_map(chan_idx).right.end()){
	  for(right_const_iterator ir=targs.first;ir!=targs.second;ir++){
	    
	    transition_probs(num_of_transition)=neg_pjump(t,dt,chan_idx,ens_it,ir);

	    target_states(num_of_transition)=ir->second;
	    
	    num_of_transition += 1;
	  }
	}
      }
    }
    //if num == 0 then there is no transitions from state hh
    if(num_of_transition != 0){
      transition_probs.resize(num_of_transition,true);
      target_states.resize(num_of_transition,true);
      //This also updates the ensemble
      make_jumps();
    }
  }
}

double nmqj_alg::pos_pjump(double t, double dt,int chan_idx,
			   std::map<int,int>::iterator it){
  rate delta=channels::get_decay_rate(chan_idx);
  return dt*delta(t)*
	      pow(norm_2(prod(channels::get_op(chan_idx),
			      effective_ensemble::get_state(it->first))),2);
}


double nmqj_alg::neg_pjump(double t,double dt, int chan_idx,
			   std::map<int,int>::iterator it,right_const_iterator ir){
  rate delta=channels::get_decay_rate(chan_idx);
  return dt*std::abs(delta(t))*
	      effective_ensemble::P(ir->second)/
	      effective_ensemble::P(it->first)*
	      pow(norm_2(prod(channels::get_op(chan_idx),
			      effective_ensemble::get_state(ir->second))),2);
}




